package fem2;

import fem2.enu.MassMatrixType;
import inf.jlinalg.Array2DMatrix;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;

import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class SequentialAssembler extends Assembler {

	public SequentialAssembler(Model m) {
		super(m);
	}

	@Override
	public void assembleKt(IMatrix kt) {
		Model m = getModel();
		/*
		 * get the array of dimension of element stiffness matrix
		 */
		Set<Integer> dimSet = new HashSet<Integer>();
		for (int i = 0; i < m.countElements(); i++) {
			dimSet.add(m.getElement(i).getMeshPart().countDofs());
		}
		/*
		 * create a map from dimension to empty matrix
		 */
		Map<Integer, IMatrix> mapKte = new HashMap<Integer, IMatrix>();
		for (int d : dimSet) {
			IMatrix dummy = new Array2DMatrix(d, d);
			mapKte.put(d, dummy);
		}
		/*
		 * assembly
		 */
		for (int i = 0; i < m.countElements(); i++) {
			Element e = m.getElement(i);
			if (!e.isInactive()) {
				IMatrix ke = mapKte.get(e.getMeshPart().countDofs());
				BLAM.zero(ke);
				ke = (IMatrix) e.computeKte(ke);

				int[] dofs = e.getMeshPart().getDofNumbers();

				for (int j = 0; j < ke.getRowCount(); j++) {
					for (int k = 0; k < ke.getColumnCount(); k++) {
						if (dofs[j] != -1 && dofs[k] != -1) {
							kt.add(dofs[j], dofs[k], ke.get(j, k));
						}
					}
				}
			}
		}
	}

	@Override
	public void assembleRi(double[] ri) {
		Model m = getModel();
		/*
		 * get the array of dimension of element stiffness matrix
		 */
		Set<Integer> dimSet = new HashSet<Integer>();
		for (int i = 0; i < m.countElements(); i++) {
			dimSet.add(m.getElement(i).getMeshPart().countDofs());
		}
		/*
		 * create a map from dimension to empty vector
		 */
		Map<Integer, IVector> mapRie = new HashMap<Integer, IVector>();
		for (int d : dimSet) {
			IVector dummy = new ArrayVector(d);
			mapRie.put(d, dummy);
		}
		/*
		 * assembly
		 */
		for (int i = 0; i < m.countElements(); i++) {
			Element e = m.getElement(i);
			if (!e.isInactive()) {
				IVector rie = mapRie.get(e.getMeshPart().countDofs());
				Arrays.fill(((ArrayVector) rie).getValues(), 0.0);
				rie = (IVector) e.computeRie(rie);

				int[] dofs = e.getMeshPart().getDofNumbers();
				for (int j = 0; j < dofs.length; j++) {
					if (dofs[j] != -1) {
						ri[dofs[j]] += rie.get(j);
					}
				}
			}
		}
	}

	@Override
	public void assembleM(IMatrix m, MassMatrixType massType) {
		Model M = getModel();
		/*
		 * get the array of dimension of element stiffness matrix
		 */
		Set<Integer> dimSet = new HashSet<Integer>();
		for (int i = 0; i < M.countElements(); i++) {
			dimSet.add(M.getElement(i).getMeshPart().countDofs());
		}
		/*
		 * create a map from dimension to empty matrix
		 */
		Map<Integer, IMatrix> mapMe = new HashMap<Integer, IMatrix>();
		for (int d : dimSet) {
			IMatrix dummy = new Array2DMatrix(d, d);
			mapMe.put(d, dummy);
		}
		/*
		 * assembly
		 */
		for (int i = 0; i < M.countElements(); i++) {
			Element e = M.getElement(i);
			if (!e.isInactive()) {
				IMatrix me = mapMe.get(e.getMeshPart().countDofs());
				BLAM.zero(me);
				me = (IMatrix) e.computeMe(me, massType);

				int[] dofs = e.getMeshPart().getDofNumbers();

				for (int j = 0; j < me.getRowCount(); j++) {
					for (int k = 0; k < me.getColumnCount(); k++) {
						if (dofs[j] != -1 && dofs[k] != -1) {
							m.add(dofs[j], dofs[k], me.get(j, k));
						}
					}
				}
			}
		}
	}

}
